package svc

import (
	"gitee.com/dn-jinmin/chatcms/internal/config"
	"gitee.com/dn-jinmin/chatcms/internal/model"
	"gitee.com/dn-jinmin/chatcms/pkg/mongox" // 添加 mongox 包导入
	milvus "github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/tmc/langchaingo/callbacks"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/ollama"
)

type ServiceContext struct {
	Config     config.Config
	Embedder   *ollama.LLM
	LLm        llms.Model
	DefaultLLM llms.Model
	Callback   callbacks.Handler
	model.ConversineModel
	model.SpeechcraftModel
}

func NewServiceContext(c config.Config) (*ServiceContext, error) {
	mongoDb, err := mongox.MongodbDatabase(&mongox.MongodbConfig{
		Password: c.Mongo.Password,
		Hosts:    c.Mongo.Hosts,
		Port:     c.Mongo.Port,
		Database: c.Mongo.Database,
		Params:   c.Mongo.Params,
		User:     c.Mongo.User,
	})
	if err != nil {
		return nil, err
	}

	embedder, err := ollama.New(ollama.WithModel(c.Ollamas.Embedder))
	if err != nil {
		return nil, err
	}
	embedder.CallbacksHandler = new(callbacks.LogHandler)

	llm, err := ollama.New()
	if err != nil {
		return nil, err
	}
	llm.CallbacksHandler = new(callbacks.LogHandler)

	defaultLLM, err := ollama.New(ollama.WithModel(c.Ollamas.Default))
	if err != nil {
		return nil, err
	}
	defaultLLM.CallbacksHandler = new(callbacks.LogHandler)

	return &ServiceContext{
		Config:           c,
		LLm:              llm,
		DefaultLLM:       defaultLLM,
		Embedder:         embedder,
		Callback:         callbacks.LogHandler{},
		ConversineModel:  model.NewConversineModel(mongoDb),
		SpeechcraftModel: model.NewSpeechcraftModel(mongoDb, embedder, milvus.Config{Address: c.Milvus.Uri}, c.SpeechcraftOffsetPath),
	}, nil
}
